package handlers

import (
	"context"
	"crypto/md5"
	"fmt"
	"generalhandler/conf"
	"generalhandler/utils"
	"github.com/bitly/go-simplejson"
	"github.com/wangthomas/gobloomfield/client"
	"regexp"
	"strings"
	"time"
)

type UrlUniqHandler struct {
	bloomClient client.Client
}

func UrlUniqHandlerInstance() *UrlUniqHandler {
	instance := new(UrlUniqHandler)
	var err error
	instance.bloomClient, err = client.NewBloomClient(conf.UrlUniqConfigInstance().BloomSvrHost,
		time.Duration(conf.UrlUniqConfigInstance().Timeout) * time.Second)
	if nil != err {
		panic(fmt.Sprintf("failed to NewBloomClient, cuz %s", err.Error()))
	}
	return instance
}

func (uuh *UrlUniqHandler) Handle(bytesMsg []byte, log Logger, pub Publisher) error {

	start := time.Now().Unix()
	msg, err := simplejson.NewJson(bytesMsg)
	if nil != err {
		log.Error("failed to marshal %s, cuz %s", string(bytesMsg), err.Error())
		return nil
	}

	strDomain, err := msg.Get("domain").String()
	if nil != err {
		log.Warn("found no domain in %s", string(bytesMsg))
		return nil
	}
	strTkId, err := msg.Get("tkid").String()
	if nil != err {
		log.Warn("found no tkid in %s", string(bytesMsg))
		return nil
	}
	strSubDomain, err := msg.Get("subdomain").String()
	if nil != err {
		log.Warn("found no subdomain in %s", string(bytesMsg))
		return nil
	}
	strPrefix, _ := msg.Get("prefix").String()
	strPort, _ := msg.Get("port").String()

	strIgnore, _ := msg.Get("ignore").String()
	var ignoreMap map[string]interface{}
	ignoreMap = make(map[string]interface{})
	utils.StringToMap(strIgnore,ignoreMap)

	list, err := msg.Get("scan_result").Get("urls").Array()
	if nil != err {
		log.Error("found no urls in %s", string(bytesMsg))
		return nil
	}
	re, _ := regexp.Compile("(=[^&|^$]+)")
	ctx := context.Background()
	var urls []string
	var keys []string
	for _, e := range list {
		url, ok := e.(string)
		if !ok {
			log.Warn("unexcept data %v", e)
			continue
		}
		// skip the url without parameters
		if !strings.Contains(url, "=") {
			continue
		}
		key := fmt.Sprintf("%x", md5.Sum(re.ReplaceAll([]byte(url), []byte(""))))
		urls = append(urls, e.(string))
		keys = append(keys, key)
	}
	var hasKeys []bool
	if 0 != len(urls) {
		hasKeys, err = uuh.bloomClient.Add(ctx, strTkId, keys)
		if nil != err {
			log.Error("failed to bloom add %v, cuz %s", keys, err.Error())
			return nil
		}
		for i, hasKey := range hasKeys {
			if hasKey {
				continue
			}
			msgNew := simplejson.New()
			msgNew.Set("start", start)
			msgNew.Set("domain", strDomain)
			msgNew.Set("tkid", strTkId)
			msgNew.Set("subdomain", strSubDomain)
			if 0 != len(strPrefix) {
				msgNew.Set("prefix", strPrefix)
			}
			if 0 != len(strPort) {
				msgNew.Set("port", strPort)
			}
			msgNew.Set("url", urls[i])
			msgNew.Set("end", time.Now().Unix())
			bytesNew, err := msgNew.MarshalJSON()
			if nil != err {
				log.Error("failed to marshal %v, cuz %s", msgNew, err.Error())
				return nil
			}
			for _, handler := range conf.MQConfigInstance().NextHandlers {
				if _, ok := ignoreMap[handler]; ok{
					continue
				}
				pub.PublishMessage(handler, bytesNew)
			}
		}
	}

	return nil
}

